from datetime import datetime
import pandas as pd

from gm.api import *

myquant_token = "3f4d2641277cdf9fb88338ccafd61b1bbd20cbe9"
set_token(myquant_token)
symbol = "rb99"
exchange = f"SHFE"
start_date = datetime(year=2020, month=11, day=1, hour=0, minute=0, second=0)
target_sec_id = symbol[:2] if symbol[1].isalpha() else symbol[0]

contract_list = []
price_tick = 0
ins = get_instruments(symbols=None, exchanges=exchange, sec_types=[4], names=None, skip_suspended=True, skip_st=True, fields="sec_id, trade_date, price_tick", df=False)
for c in ins:
    sec_id = c["sec_id"][:2] if c["sec_id"][1].isalpha() else c["sec_id"][0]
    dt = datetime(year=c["trade_date"].year, month=c["trade_date"].month, day=c["trade_date"].day)
    if sec_id == target_sec_id and dt >= start_date:
        if not price_tick:
            price_tick = c["price_tick"]
        contract_list.append(c["sec_id"])

weight_average_cols = ["open", "high", "low", "close"]
weight_sum_cols = ["volume"]
all_interest_dict = {}
all_price_data = {}
for contract_id in contract_list:
    symbols = exchange + "." + contract_id
    df = history(symbol=symbols, frequency='60s', start_time=start_date.strftime("%Y-%m-%d %H:%M:%S"),
                 end_time=datetime.now().strftime("%Y-%m-%d %H:%M:%S"), df=True, fill_missing="Last")

    order_price_data = df
    if order_price_data is None:
        continue
    order_price_data.set_index("bob", inplace=True)
    #order_price_data = order_price_data.truncate(before=start_date)
    all_interest_dict[contract_id] = order_price_data["position"]
    all_price_data[contract_id] = order_price_data
    order_price_data.to_csv(f"D:\\learn_and_test_data\\{contract_id}_1m.csv")
print("date collected")
all_interest_df = pd.DataFrame(all_interest_dict)
# 计算总持仓和每个合约的持仓比例
all_interest_df.fillna(method="ffill", inplace=True)
all_interest_df["position"] = all_interest_df.sum(axis=1)
for order_id in all_interest_dict.keys():
    all_interest_df[order_id] = all_interest_df[order_id] / all_interest_df["position"]
all_interest_df.to_csv(f"D:\\learn_and_test_data\\rb99_position_1m.csv")
index_data = pd.DataFrame(index=all_interest_df.index)
index_data["symbol_id"] = symbol

for col in weight_average_cols:
    index_data[col] = 0
    for order_id, price_data in all_price_data.items():
        index_data[col] += (price_data[col] * all_interest_df[order_id]).fillna(method="ffill").fillna(0)
    a = index_data[col] / price_tick
    b = a.round(0)
    c = b * price_tick
    index_data[col] = c
for col in weight_sum_cols:
    index_data[col] = 0
    for order_id, price_data in all_price_data.items():
        buff = index_data[col] + price_data[col]
        index_data[col] += buff.fillna(0)
index_data["position"] = all_interest_df["position"]

index_data.to_csv("D:\\learn_and_test_data\\rb99_1m.csv")
